#  ==============================导入相关包==============================
import random
import os

import numpy as np
import torch
from torch import nn

import logging
LOGGER = logging.getLogger(__name__)
LOGGER.addHandler(logging.NullHandler())

logging.basicConfig(level=logging.INFO, format='%(asctime)-15s %(levelname)s: %(message)s')

# set seed
seed = 666
random.seed(seed)
np.random.seed(seed)
torch.cuda.manual_seed(seed)
torch.manual_seed(seed)

#  ==============================设置是否需要使用显卡==============================
# set cuda
gpu = 0
use_cuda = gpu >= 0 and torch.cuda.is_available()
if use_cuda:
    torch.cuda.set_device(gpu)
    device = torch.device("cuda", gpu)
else:
    device = torch.device("cpu")
logging.info("Use cuda: %s, gpu id: %d.", use_cuda, gpu)

#  ==============================将数据集划分成10个fold==============================
# split data to 10 fold
fold_num = 10
data_file = '../data/train_set.csv'
import pandas as pd


def all_data2fold(fold_num, num=10000):
    """将所有文档分成fold_num折"""
    fold_data = []
    f = pd.read_csv(data_file, sep='\t', encoding='UTF-8')
    texts = f['text'].tolist()[:num]
    labels = f['label'].tolist()[:num]

    total = len(labels)

    index = list(range(total))
    np.random.shuffle(index)

    all_texts = []
    all_labels = []
    for i in index:
        all_texts.append(texts[i])
        all_labels.append(labels[i])

    label2id = {}
    for i in range(total):
        label = str(all_labels[i])
        if label not in label2id:
            label2id[label] = [i]
        else:
            label2id[label].append(i)

    all_index = [[] for _ in range(fold_num)]
    for label, data in label2id.items():
        # print(label, len(data))
        batch_size = int(len(data) / fold_num)
        other = len(data) - batch_size * fold_num
        for i in range(fold_num):
            cur_batch_size = batch_size + 1 if i < other else batch_size
            # print(cur_batch_size)
            batch_data = [data[i * batch_size + b] for b in range(cur_batch_size)]
            all_index[i].extend(batch_data)

    batch_size = int(total / fold_num)
    other_texts = []
    other_labels = []
    other_num = 0
    start = 0
    for fold in range(fold_num):
        num = len(all_index[fold])
        texts = [all_texts[i] for i in all_index[fold]]
        labels = [all_labels[i] for i in all_index[fold]]

        if num > batch_size:
            fold_texts = texts[:batch_size]
            other_texts.extend(texts[batch_size:])
            fold_labels = labels[:batch_size]
            other_labels.extend(labels[batch_size:])
            other_num += num - batch_size
        elif num < batch_size:
            end = start + batch_size - num
            fold_texts = texts + other_texts[start: end]
            fold_labels = labels + other_labels[start: end]
            start = end
        else:
            fold_texts = texts
            fold_labels = labels

        assert batch_size == len(fold_labels)

        # shuffle
        index = list(range(batch_size))
        np.random.shuffle(index)

        shuffle_fold_texts = []
        shuffle_fold_labels = []
        for i in index:
            shuffle_fold_texts.append(fold_texts[i])
            shuffle_fold_labels.append(fold_labels[i])

        data = {'label': shuffle_fold_labels, 'text': shuffle_fold_texts}
        fold_data.append(data)

    logging.info("Fold lens %s", str([len(data['label']) for data in fold_data]))

    return fold_data


fold_data = all_data2fold(10, num=99999999999)
#  fold_data = all_data2fold(10)

# build train, dev, test data
fold_id = 9

# dev
dev_data = fold_data[fold_id]

# train
train_texts = []
train_labels = []
for i in range(0, fold_id):
    data = fold_data[i]
    train_texts.extend(data['text'])
    train_labels.extend(data['label'])

train_data = {'label': train_labels, 'text': train_texts}

# test
test_data_file = '../data/test_a.csv'
f = pd.read_csv(test_data_file, sep='\t', encoding='UTF-8')
texts = f['text'].tolist()
test_data = {'label': [0] * len(texts), 'text': texts}

#  ==============================建立vocab==============================
# build vocab
from collections import Counter
from transformers import BasicTokenizer

basic_tokenizer = BasicTokenizer()


class Vocab():
    def __init__(self, train_data, mode='basic'):
        self.min_count = 5
        self.pad = 0
        self.unk = 1
        if mode == 'bert':
            self._id2word = ['[PAD]', '[UNK]', '[CLS]', '[SEP]', '[MASK]'] #  list
            self._id2extword = ['[PAD]', '[UNK]', '[CLS]', '[SEP]', '[MASK]']
        else:
            self._id2word = ['[PAD]', '[UNK]'] #  list
            self._id2extword = ['[PAD]', '[UNK]']

        self._id2label = []
        self.target_names = []

        self.build_vocab(train_data)

        reverse = lambda x: dict(zip(x, range(len(x))))
        self._word2id = reverse(self._id2word) #  dict
        self._label2id = reverse(self._id2label) #  dict

        logging.info("Build vocab: words %d, labels %d." % (self.word_size, self.label_size))

    def build_vocab(self, data):
        self.word_counter = Counter()

        for text in data['text']:
            words = text.split()
            for word in words:
                self.word_counter[word] += 1

        for word, count in self.word_counter.most_common():
            if count >= self.min_count:
                self._id2word.append(word)

        label2name = {0: '科技', 1: '股票', 2: '体育', 3: '娱乐', 4: '时政', 5: '社会', 6: '教育', 7: '财经',
                      8: '家居', 9: '游戏', 10: '房产', 11: '时尚', 12: '彩票', 13: '星座'}

        self.label_counter = Counter(data['label'])

        for label in range(len(self.label_counter)):
            count = self.label_counter[label]
            self._id2label.append(label)
            self.target_names.append(label2name[label])

    def load_pretrained_embs(self, embfile):
        """加载预训练的embedding文件，生成_id2extword以及embedding层"""
        with open(embfile, encoding='utf-8') as f:
            lines = f.readlines()
            items = lines[0].split()
            word_count, embedding_dim = int(items[0]), int(items[1])

        index = len(self._id2extword) #  预训练embedding的基础上还需要加上['[PAD]', '[UNK]']
        embeddings = np.zeros((word_count + index, embedding_dim))
        for line in lines[1:]:
            values = line.split()
            self._id2extword.append(values[0])
            vector = np.array(values[1:], dtype='float64')
            embeddings[self.unk] += vector #  这里在UNK对应的行上将所有词的embedding向量相加
            embeddings[index] = vector
            index += 1

        embeddings[self.unk] = embeddings[self.unk] / word_count #  将累加的所有词的embedding向量除以单词数
        embeddings = embeddings / np.std(embeddings)

        reverse = lambda x: dict(zip(x, range(len(x))))
        self._extword2id = reverse(self._id2extword)

        assert len(set(self._id2extword)) == len(self._id2extword)

        return embeddings

    def word2id(self, xs):
        if isinstance(xs, list):
            return [self._word2id.get(x, self.unk) for x in xs]
        return self._word2id.get(xs, self.unk)

    def extword2id(self, xs):
        if isinstance(xs, list):
            return [self._extword2id.get(x, self.unk) for x in xs]
        return self._extword2id.get(xs, self.unk)

    def label2id(self, xs):
        if isinstance(xs, list):
            return [self._label2id.get(x, self.unk) for x in xs]
        return self._label2id.get(xs, self.unk)

    def save_vocab(self, filepath):
        with open(filepath, 'w') as f:
            f.writelines([word+'\n' for word in self._id2word])

    @property
    def word_size(self):
        return len(self._id2word)

    @property
    def extword_size(self):
        return len(self._id2extword)

    @property
    def label_size(self):
        return len(self._id2label)


vocab = Vocab(train_data, 'bert')

#  ==============================创建dataset==============================
# build dataset
def sentence_split(text, vocab, max_sent_len=256, max_segment=16):
    """将一个文档拆成多个句子
    @param max_sent_len: 每个句子最多多少个单词
    @type  max_sent_len: int

    @param max_segment: 每个文档最多拆成多少个句子
    @type  max_segment: int

    @return: [[sent_len, [word0, ...]], ...]，[[该句子有多少个单词, [第一个单词, ...]], ...]
    @rtype :  list
    """
    words = text.strip().split() #  该文档的所有组成单词列表
    document_len = len(words) #  该文档共有多少个单词

    index = list(range(0, document_len, max_sent_len)) #  构建切片时的滑窗索引，将文档总单词个数每隔max_sent_len一个索引，下面就根据这个索引来将一个文档拆成多个句子
    index.append(document_len) #  文档的单词个数作为最后一个索引加入

    segments = []
    for i in range(len(index) - 1):
        segment = words[index[i]: index[i + 1]] #  根据索引对该文档的所有组成单词列表进行切片，滑窗取句子
        assert len(segment) > 0
        segment = [word if word in vocab._id2word else '<UNK>' for word in segment] #  如果每个词都在vocab中出现过则保留原词，否则标识成<UNK>
        segments.append([len(segment), segment])

    assert len(segments) > 0
    if len(segments) > max_segment: #  如果该文档拆分成的句子数超过了max_segment，则只保留头尾的int(max_segment / 2)个
        segment_ = int(max_segment / 2)
        return segments[:segment_] + segments[-segment_:]
    else:
        return segments


def get_examples(data, vocab, max_sent_len=256, max_segment=8, mode='basic'):
    """创建dataset，注意：因为有些文档较长，为了充分训练模型，采用将一个分档拆成多个句子的方式
    @param data: data = {'label': shuffle_fold_labels, 'text': shuffle_fold_texts}
    @type  data: dict

    @param max_sent_len: 每个句子最多多少个单词
    @type  max_sent_len: int

    @param max_segment: 每个文档最多拆成多少个句子
    @type  max_segment: int

    @return: [[label, doc_len, [[sent_len, [sent_id0, ...], [sent_id1, ...]], ...]]，[标签, 该文档有多少个句子, [[该句子有多少个单词, [自建embedding的id0, ...], [预训练embedding的id1, ...]], ...]]
    @rtype :  list
    """

    label2id = vocab.label2id
    examples = []
    tokenizer = WhitespaceTokenizer(vocab) if mode == 'bert' else None

    for text, label in zip(data['text'], data['label']): #  逐行遍历，取出一个文档和其对应label，注意：还将一个文档拆成了多个句子
        # label
        id = label2id(label)

        # words
        if mode == 'bert':
            sents_words = sentence_split(text, vocab, max_sent_len-2, max_segment)
            doc = []
            #  训练bert时，因为在'[PAD]', '[UNK]'基础上额外多出了'[CLS]', '[SEP]', '[MASK]'字符，所以和其他模型的vocab不一样，需要使用WhitespaceTokenizer来给句子添加上'[CLS]'和'[SEP]'再用相应的vocab进行tokenize
            for sent_len, sent_words in sents_words:
                token_ids = tokenizer.tokenize(sent_words)
                sent_len = len(token_ids)
                token_type_ids = [0] * sent_len
                doc.append([sent_len, token_ids, token_type_ids])
        else:
            sents_words = sentence_split(text, vocab, max_sent_len, max_segment) # 将一个文档拆成多个句子
            doc = []
            for sent_len, sent_words in sents_words:
                word_ids = vocab.word2id(sent_words) #  word2id，自己新建的embedding层用到的id
                extword_ids = vocab.extword2id(sent_words) # extword2id，预训练的embedding层用到的id，比如使用word2vec训练好的
                doc.append([sent_len, word_ids, extword_ids])
        examples.append([id, len(doc), doc])

    logging.info('Total %d docs.' % len(examples))
    return examples


class WhitespaceTokenizer():
    """
    主要功能：
    1. 加载vocab
    2. 给每个句子头位分别加上[CLS]和[SEP]
    3. 使用vocab对句子进行tokenize--word2id
    WhitespaceTokenizer with vocab.
    """

    def __init__(self, vocab):
        if isinstance(vocab, str):
            self._word2id = self.load_vocab(vocab)
        else:
            self._word2id = vocab._word2id
        self._id2word = {v: k for k, v in self._word2id.items()}
        self.max_len = 256
        self.unk = 1

        logging.info("Build Bert vocab with size %d." % (self.vocab_size))

    def load_vocab(self, vocab_file):
        f = open(vocab_file, 'r')
        lines = f.readlines()
        lines = list(map(lambda x: x.strip(), lines))
        word2id = dict(zip(lines, range(len(lines))))
        return word2id

    def tokenize(self, tokens):
        assert len(tokens) <= self.max_len - 2
        tokens = ["[CLS]"] + tokens + ["[SEP]"]
        output_tokens = self.token2id(tokens)
        return output_tokens

    def token2id(self, xs):
        if isinstance(xs, list):
            return [self._word2id.get(x, self.unk) for x in xs]
        return self._word2id.get(xs, self.unk)

    @property
    def vocab_size(self):
        return len(self._id2word)


def get_bert_examples(data, word_encoder, vocab, max_sent_len=256, max_segment=8):
    """
    """
    label2id = vocab.label2id
    examples = []

    for text, label in zip(data['text'], data['label']):
        # label
        id = label2id(label)

        # words
        sents_words = sentence_split(text, vocab, max_sent_len-2, max_segment)
        doc = []
        for sent_len, sent_words in sents_words:
            token_ids = word_encoder.encode(sent_words)
            sent_len = len(token_ids)
            token_type_ids = [0] * sent_len
            doc.append([sent_len, token_ids, token_type_ids])
        examples.append([id, len(doc), doc])

    logging.info('Total %d docs.' % len(examples))
    return examples

#  ==============================创建loader==============================
# build loader

def batch_slice(data, batch_size):
    batch_num = int(np.ceil(len(data) / float(batch_size)))
    for i in range(batch_num):
        cur_batch_size = batch_size if i < batch_num - 1 else len(data) - batch_size * i
        docs = [data[i * batch_size + b] for b in range(cur_batch_size)]

        yield docs


def data_iter(data, batch_size, shuffle=True, noise=1.0):
    """
    randomly permute data, then sort by source length, and partition into batches
    ensure that the length of  sentences in each batch

    将数据集拆成一个个batch，关键点：每个batch内的句子按句子长度排序
    @data: [[label, doc_len, [[sent_len, [sent_id0, ...], [sent_id1, ...]], ...]]，[标签, 该文档有多少个句子, [[该句子有多少个单词, [自建embedding的id0, ...], [预训练embedding的id1, ...]], ...]]
    @rtype :  list
    """

    batched_data = []
    if shuffle:
        np.random.shuffle(data)

        lengths = [example[1] for example in data] #  取出每个文档的组成句子数
        noisy_lengths = [- (l + np.random.uniform(- noise, noise)) for l in lengths] #  句子数上加个均匀分布的小噪声，再取负
        sorted_indices = np.argsort(noisy_lengths).tolist() #  按从小到大排列，因为上面取了负数，其实相当于按句子长度从大到小排列
        sorted_data = [data[i] for i in sorted_indices]
    else:
        sorted_data = data

    batched_data.extend(list(batch_slice(sorted_data, batch_size)))

    if shuffle:
        np.random.shuffle(batched_data)

    for batch in batched_data:
        yield batch

#  ==============================创建优化器==============================
# build optimizer
learning_rate = 2e-4
bert_lr = 5e-5
decay = .75
decay_step = 1000
from transformers import AdamW, get_linear_schedule_with_warmup


class Optimizer:
    def __init__(self, model_parameters, steps):
        self.all_params = []
        self.optims = []
        self.schedulers = []

        for name, parameters in model_parameters.items():
            if name.startswith("basic"):
                optim = torch.optim.Adam(parameters, lr=learning_rate)
                self.optims.append(optim)

                l = lambda step: decay ** (step // decay_step)
                scheduler = torch.optim.lr_scheduler.LambdaLR(optim, lr_lambda=l)
                self.schedulers.append(scheduler)
                self.all_params.extend(parameters)
            elif name.startswith("bert"):
                optim_bert = AdamW(parameters, bert_lr, eps=1e-8)
                self.optims.append(optim_bert)

                scheduler_bert = get_linear_schedule_with_warmup(optim_bert, 0, steps)
                self.schedulers.append(scheduler_bert)

                for group in parameters:
                    for p in group['params']:
                        self.all_params.append(p)
            else:
                Exception("no nameed parameters.")

        self.num = len(self.optims)

    def step(self):
        for optim, scheduler in zip(self.optims, self.schedulers):
            optim.step()
            scheduler.step()
            optim.zero_grad()

    def zero_grad(self):
        for optim in self.optims:
            optim.zero_grad()

    def get_lr(self):
        lrs = tuple(map(lambda x: x.get_lr()[-1], self.schedulers))
        lr = ' %.5f' * self.num
        res = lr % lrs
        return res

#  ==============================计算得分==============================
# some function
from sklearn.metrics import f1_score, precision_score, recall_score


def get_score(y_ture, y_pred):
    y_ture = np.array(y_ture)
    y_pred = np.array(y_pred)
    f1 = f1_score(y_ture, y_pred, average='macro') * 100
    p = precision_score(y_ture, y_pred, average='macro') * 100
    r = recall_score(y_ture, y_pred, average='macro') * 100

    return str((reformat(p, 2), reformat(r, 2), reformat(f1, 2))), reformat(f1, 2)


def reformat(num, n):
    return float(format(num, '0.' + str(n) + 'f'))

#  ==============================创建训练器==============================
# build trainer

import time
from sklearn.metrics import classification_report

clip = 5.0
epochs = 99
early_stops = 5
log_interval = 50

test_batch_size = 128
train_batch_size = 128

class Trainer():
    def __init__(self, model, vocab, save_model, save_test='sub.csv', mode='basic'):
        if os.path.exists(save_model):
            self.model = torch.load(save_model)
        else:
            self.model = model
        self.report = True

        if use_cuda:
            self.model.to(device)

        self.save_model = save_model
        self.save_test = save_test

        self.train_data = get_examples(train_data, vocab, mode=mode)
        self.batch_num = int(np.ceil(len(self.train_data) / float(train_batch_size)))
        self.dev_data = get_examples(dev_data, vocab, mode=mode)
        self.test_data = get_examples(test_data, vocab, mode=mode)

        # criterion
        self.criterion = nn.CrossEntropyLoss()

        # label name
        self.target_names = vocab.target_names

        # optimizer
        self.optimizer = Optimizer(model.all_parameters, steps=self.batch_num * epochs)

        # count
        self.step = 0
        self.early_stop = -1
        self.best_train_f1, self.best_dev_f1 = 0, 0
        self.last_epoch = epochs

    def train(self):
        logging.info('Start training...')
        for epoch in range(1, epochs + 1):
            train_f1 = self._train(epoch)

            dev_f1 = self._eval(epoch)

            if self.best_dev_f1 <= dev_f1:
                logging.info(
                    "Exceed history dev = %.2f, current dev = %.2f" % (self.best_dev_f1, dev_f1))
                torch.save(self.model.state_dict(), self.save_model)

                self.best_train_f1 = train_f1
                self.best_dev_f1 = dev_f1
                self.early_stop = 0
            else:
                self.early_stop += 1
                if self.early_stop == early_stops:
                    logging.info(
                        "Eearly stop in epoch %d, best train: %.2f, dev: %.2f" % (
                            epoch - early_stops, self.best_train_f1, self.best_dev_f1))
                    self.last_epoch = epoch
                    break

    def test(self):
        self.model.load_state_dict(torch.load(self.save_model))
        self._eval(self.last_epoch + 1, test=True)

    def _train(self, epoch):
        self.optimizer.zero_grad()
        self.model.train()

        start_time = time.time()
        epoch_start_time = time.time()
        overall_losses = 0
        losses = 0
        batch_idx = 1
        y_pred = []
        y_true = []
        for batch_data in data_iter(self.train_data, train_batch_size, shuffle=True):
            torch.cuda.empty_cache()
            batch_inputs, batch_labels = self.batch2tensor(batch_data)
            batch_outputs = self.model(batch_inputs)
            loss = self.criterion(batch_outputs, batch_labels)
            loss.backward()

            loss_value = loss.detach().cpu().item()
            losses += loss_value
            overall_losses += loss_value

            y_pred.extend(torch.max(batch_outputs, dim=1)[1].cpu().numpy().tolist())
            y_true.extend(batch_labels.cpu().numpy().tolist())

            nn.utils.clip_grad_norm_(self.optimizer.all_params, max_norm=clip)
            for optimizer, scheduler in zip(self.optimizer.optims, self.optimizer.schedulers):
                optimizer.step()
                scheduler.step()
            self.optimizer.zero_grad()

            self.step += 1

            if batch_idx % log_interval == 0:
                elapsed = time.time() - start_time

                lrs = self.optimizer.get_lr()
                logging.info(
                    '| epoch {:3d} | step {:3d} | batch {:3d}/{:3d} | lr{} | loss {:.4f} | s/batch {:.2f}'.format(
                        epoch, self.step, batch_idx, self.batch_num, lrs,
                        losses / log_interval,
                        elapsed / log_interval))

                losses = 0
                start_time = time.time()

            batch_idx += 1

        overall_losses /= self.batch_num
        during_time = time.time() - epoch_start_time

        # reformat
        overall_losses = reformat(overall_losses, 4)
        score, f1 = get_score(y_true, y_pred)

        logging.info(
            '| epoch {:3d} | score {} | f1 {} | loss {:.4f} | time {:.2f}'.format(epoch, score, f1,
                                                                                  overall_losses,
                                                                                  during_time))
        if set(y_true) == set(y_pred) and self.report:
            report = classification_report(y_true, y_pred, digits=4, target_names=self.target_names)
            logging.info('\n' + report)

        return f1

    def _eval(self, epoch, test=False):
        self.model.eval()
        start_time = time.time()
        data = self.test_data if test else self.dev_data
        y_pred = []
        y_true = []
        with torch.no_grad():
            for batch_data in data_iter(data, test_batch_size, shuffle=False):
                torch.cuda.empty_cache()
                batch_inputs, batch_labels = self.batch2tensor(batch_data)
                batch_outputs = self.model(batch_inputs)
                y_pred.extend(torch.max(batch_outputs, dim=1)[1].cpu().numpy().tolist())
                y_true.extend(batch_labels.cpu().numpy().tolist())

            score, f1 = get_score(y_true, y_pred)

            during_time = time.time() - start_time

            if test:
                df = pd.DataFrame({'label': y_pred})
                df.to_csv(self.save_test, index=False, sep=',')
            else:
                logging.info(
                    '| epoch {:3d} | dev | score {} | f1 {} | time {:.2f}'.format(epoch, score, f1,
                                                                              during_time))
                if set(y_true) == set(y_pred) and self.report:
                    report = classification_report(y_true, y_pred, digits=4, target_names=self.target_names)
                    logging.info('\n' + report)

        return f1

    def batch2tensor(self, batch_data):
        '''
            [[label, doc_len, [[sent_len, [sent_id0, ...], [sent_id1, ...]], ...]]
        '''
        batch_size = len(batch_data)
        doc_labels = []
        doc_lens = []
        doc_max_sent_len = []
        for doc_data in batch_data:
            doc_labels.append(doc_data[0])
            doc_lens.append(doc_data[1])
            sent_lens = [sent_data[0] for sent_data in doc_data[2]]
            max_sent_len = max(sent_lens)
            doc_max_sent_len.append(max_sent_len)

        max_doc_len = max(doc_lens) #  取包含最多句子数的文档的句子数
        max_sent_len = max(doc_max_sent_len) #  取包含最多单词数的句子的单词数

        batch_inputs1 = torch.zeros((batch_size, max_doc_len, max_sent_len), dtype=torch.int64)
        batch_inputs2 = torch.zeros((batch_size, max_doc_len, max_sent_len), dtype=torch.int64)
        batch_masks = torch.zeros((batch_size, max_doc_len, max_sent_len), dtype=torch.float32)
        batch_labels = torch.LongTensor(doc_labels)

        for b in range(batch_size):
            for seg_idx in range(doc_lens[b]):
                sent_data = batch_data[b][2][seg_idx]
                for word_idx in range(sent_data[0]):
                    batch_inputs1[b, seg_idx, word_idx] = sent_data[1][word_idx]
                    batch_inputs2[b, seg_idx, word_idx] = sent_data[2][word_idx]
                    batch_masks[b, seg_idx, word_idx] = 1

        if use_cuda:
            batch_inputs1 = batch_inputs1.to(device)
            batch_inputs2 = batch_inputs2.to(device)
            batch_masks = batch_masks.to(device)
            batch_labels = batch_labels.to(device)

        return (batch_inputs1, batch_inputs2, batch_masks), batch_labels

#  ==============================训练和测试==============================
#  from models import textcnn, textrnn, bert
#  model = textcnn.Model(vocab)

# train
#  save_model = './cnn.bin'
#  save_test = './cnn.csv'

save_model = './rnn.bin'
save_test = './rnn.csv'

#  save_model = './bert.bin'
#  save_test = './bert.csv'

model = Model(vocab)
trainer = Trainer(model, vocab, save_model, save_test, 'bert')
trainer.train()

# test
trainer.test()
